﻿#include "imghelper.h"
#include <QFile>
#include <QDebug>
#include <QImage>

QByteArray imgHelper::getImgFromFile(const QString &in_imgPath)
{
    QFile imgFile(in_imgPath);
    if(!imgFile.exists() || !imgFile.open(QIODevice::ReadOnly))
        return "";
    QByteArray imgArr = imgFile.readAll();
    imgFile.close();
    return imgArr;
}

QVector<QVector<unsigned char>> imgHelper::rgbSplit(const unsigned char *in_srcImg, int width, int height)
{
    QVector<unsigned char> rImg;
    QVector<unsigned char> gImg;
    QVector<unsigned char> bImg;

    rImg.reserve(height*width);
    gImg.reserve(height*width);
    bImg.reserve(height*width);

    for(int i = 0;i < height*width*3;i+=3)
    {
        rImg.append(in_srcImg[i]);
        gImg.append(in_srcImg[i+1]);
        bImg.append(in_srcImg[i+2]);
    }

    QVector<QVector<unsigned char>> rgbImg;
    rgbImg.reserve(3);
    rgbImg.append(std::move(rImg));
    rgbImg.append(std::move(gImg));
    rgbImg.append(std::move(bImg));
    return rgbImg;
}

QVector<QImage> imgHelper::rgbSplit(const QImage &in_img)
{
    QVector<QImage> channels;
    if (in_img.isNull()) {
        return channels;
    }

    QImage rgbImg = in_img;
    if (rgbImg.format() != QImage::Format_RGB888) {
        rgbImg = rgbImg.convertToFormat(QImage::Format_RGB888);
    }

    const int width = rgbImg.width();
    const int height = rgbImg.height();

    QImage rImg(width, height, QImage::Format_Grayscale8);
    QImage gImg(width, height, QImage::Format_Grayscale8);
    QImage bImg(width, height, QImage::Format_Grayscale8);

    for (int y = 0; y < height; y++)
    {
        unsigned char * rgbLine = rgbImg.scanLine(y);
        unsigned char * rLine = rImg.scanLine(y);
        unsigned char * gLine = gImg.scanLine(y);
        unsigned char * bLine = bImg.scanLine(y);
        for (int i = 0,index =0; i < width*3; i+=3,index++)
        {
            rLine[index] = rgbLine[i];
            gLine[index] = rgbLine[i+1];
            bLine[index] = rgbLine[i+2];
        }
    }

    channels.reserve(3);
    channels.append(std::move(rImg));
    channels.append(std::move(gImg));
    channels.append(std::move(bImg));
    return channels;
}

QVector<QVector<unsigned char> > imgHelper::yuvSplit(const unsigned char *in_srcImg, int width, int height)
{
    if(in_srcImg == nullptr || width<=0 || height<=0)
        return QVector<QVector<unsigned char> >();
    // 1. 预分配内存
    QVector<unsigned char> yImg;
    QVector<unsigned char> uImg;
    QVector<unsigned char> vImg;
    yImg.resize(width*height);
    uImg.resize(width*height/4);
    vImg.resize(width*height/4);
    // 2. yuv的起始指针
    const unsigned char * yPtr = in_srcImg;
    const unsigned char * uPtr = yPtr+width*height;
    const unsigned char * vPtr = uPtr+width*height/4;

    // 3. 复制到对应的内存中
    memcpy(yImg.data(),yPtr,width*height);
    memcpy(uImg.data(),uPtr,width*height/4);
    memcpy(vImg.data(),vPtr,width*height/4);

    QVector<QVector<unsigned char>> yuvImg;
    yuvImg.reserve(3);
    yuvImg.append(std::move(yImg));
    yuvImg.append(std::move(uImg));
    yuvImg.append(std::move(vImg));
    return yuvImg;
}

QVector<unsigned char> imgHelper::yuv420pToRGB24(const unsigned char *in_srcImg, int width, int height)
{
    if(in_srcImg==nullptr || width<=0 || height<=0)
        return QVector<unsigned char>();

    // 1.分别获得y/u/v的开始指针
    const unsigned char *yPtr = in_srcImg;
    const unsigned char *uPtr = yPtr + width*height;
    const unsigned char *vPtr = uPtr+(width/2)*(height/2);

    // 预分配RGB内存
    QVector<unsigned char> rgbImg;
    rgbImg.resize(width*height*3);

    // 2.遍历每个像素
    for(int y = 0; y < height; ++y)
    {
        for(int x = 0; x < width; ++x)
        {
            //获取坐标
            int yIdx = y * width + x;
            int uvIdx = (y/2)*(width/2)+(x/2);

            //转换成rgb数值范围
            const int Y = yPtr[yIdx]-16;
            const int U = uPtr[uvIdx]-128;
            const int V = vPtr[uvIdx]-128;

            //将yuv转rgb数值
            int R = Y + 1.402 * V;
            int G = Y - 0.344 * U - 0.714 * V;
            int B = Y + 1.772 * U;

            // 限制值在 [0, 255] 范围内
            R = std::clamp(R, 0, 255);
            G = std::clamp(G, 0, 255);
            B = std::clamp(B, 0, 255);

            //一个像素三个字节
            int rgbIdx = (y*width+x)*3;
            rgbImg[rgbIdx] = static_cast<unsigned char>(R);
            rgbImg[rgbIdx+1] = static_cast<unsigned char>(G);
            rgbImg[rgbIdx+2] = static_cast<unsigned char>(B);
        }
    }
    return rgbImg;
}

QVector<unsigned char> imgHelper::rgb24ToYUV420p(const unsigned char *in_srcImg, int width, int height)
{
    if(in_srcImg==nullptr || width<=0 || height<=0)
        return QVector<unsigned char>();

    // 1. 预分配yuv内存
    QVector<unsigned char> yuvImg;
    yuvImg.resize(width*height*1.5);

    // 2. 遍历每个像素
    for(int y = 0; y < height; ++y)
    {
        for(int x = 0; x < width; ++x)
        {
            //获取rgb数据
            int rgbIdx = (y*width+x)*3;
            unsigned char R = in_srcImg[rgbIdx];
            unsigned char G = in_srcImg[rgbIdx+1];
            unsigned char B = in_srcImg[rgbIdx+2];

            //转换成YUV数据
            int Y = (299*R + 587*G + 114*B)/1000;
            Y = std::clamp(Y, 0, 255);
            int yIndx = y*width+x;
            yuvImg[yIndx] = Y;

            //2x2 4个像素共享
            if(x%2==0 && y%2==0)
            {
                int U = (-147*R - 289*G + 436*B)/1000 + 128;
                int V = (615*R - 515*G - 100*B)/1000 + 128;
                U = std::clamp(U, 0, 255);
                V = std::clamp(V, 0, 255);
                int uIndx = (width*height)+ (y/2)*(width/2)+(x/2);
                int vIndx = (width*height) + (width/2)*(height/2) + (y/2)*(width/2)+(x/2);
                yuvImg[uIndx] = U;
                yuvImg[vIndx] = V;
            }
        }
    }
    return yuvImg;
}

void imgHelper::saveChannelsToFiles(const QVector<QVector<unsigned char> > &in_imgs, int width, int height, const QString &format)
{
    if(in_imgs.size() == 0 || width<=0 || height<=0 || in_imgs.size()!=format.size())
        return;

    for(int i=0;i<in_imgs.size();++i)
    {
        int tmpWidth = width;
        int tmpHeight = height;
        if(format[i] == 'u' || format[i] == 'v')
        {
            tmpWidth = width/2;
            tmpHeight = height/2;
        }

        QString  fileName = QString("channel%1_%2x%3.%4").arg(i+1).arg(tmpWidth).arg(tmpHeight).arg(format[i]);
        QFile file(fileName);
        if(file.open(QIODevice::WriteOnly))
        {
            file.write(reinterpret_cast<const char*>(in_imgs[i].data()),in_imgs[i].size());
            file.close();
        }
    }
}

QVector<unsigned char> imgHelper::yuvToGray(const unsigned char *in_srcImg, int width, int height)
{
    QVector<unsigned char> yuvGrayImg;
    if(in_srcImg==nullptr || width<=0 || height<=0)
        return yuvGrayImg;

    yuvGrayImg.resize(width*height*1.5);

    //将y分量直接复制
    memcpy(yuvGrayImg.data(),in_srcImg,width*height);
    //将uv分量设置为128 灰度
    std::fill_n(yuvGrayImg.data()+width*height,width*height/2,128);
    return yuvGrayImg;
}

QVector<unsigned char> imgHelper::yuvToYhalf(const unsigned char *in_srcImg, int width, int height)
{
    QVector<unsigned char> yuvHalfImg;
    if(in_srcImg==nullptr || width<=0 || height<=0)
        return yuvHalfImg;

    yuvHalfImg.resize(width*height*1.5);

    //将y分量减半
    for(int i=0;i<width*height;++i)
    {
        yuvHalfImg[i] = in_srcImg[i]/2;
    }
    //将uv分量直接复制
    memcpy(yuvHalfImg.data()+width*height,in_srcImg+width*height,width*height/2);

    return yuvHalfImg;
}

int imgHelper::psnr(const unsigned char *in_srcImg, const unsigned char *in_delImg, int width, int height)
{
    if (in_srcImg == nullptr || in_delImg == nullptr || width <= 0 || height <= 0) {
        return -1;
    }

    const int totalPixels = width * height;
    unsigned long long sumSquaredError = 0;

    // 2. 计算所有像素的平方误差和
    for (int i = 0; i < totalPixels; ++i) {
        const int srcVal = static_cast<int>(in_srcImg[i]);
        const int delVal = static_cast<int>(in_delImg[i]);
        const int diff = srcVal - delVal;
        sumSquaredError += static_cast<unsigned long long>(diff * diff);
    }

    // 3. 处理全零误差（图像完全相同）
    if (sumSquaredError == 0) {
        return 100; // 返回 100 dB 表示完美质量
    }

    // 4. 计算 MSE 和 PSNR
    const double mse = static_cast<double>(sumSquaredError) / totalPixels;
    const double psnr = 10.0 * std::log10(255.0 * 255.0 / mse);

    // 5. 四舍五入为整数返回
    return static_cast<int>(psnr + 0.5);
}
